import os.path
import time
import pandas as pd
import numpy as np
from operator_switch import pct_change
from operator_switch import stddev_bottleneck as stddev
from operator_switch import ts_argmax_bottleneck as ts_argmax
from operator_switch import rank_bottleneck as rank


data_dir = os.path.expanduser("./data")
data = pd.read_csv(os.path.join(data_dir, 'dataPerformance.csv'))

close = data.pivot(index='tradetime', columns='securityid', values='close')

close_ndarr = close.values
rets_ndarr = pct_change(close_ndarr)

t1 = time.time()

inner = close_ndarr.copy()
np.putmask(inner, rets_ndarr < 0, stddev(rets_ndarr, 20))
result_ndarr = rank(ts_argmax(inner ** 2, 5)) - 0.5

t4 = time.time()

result = pd.DataFrame(result_ndarr, columns=close.columns)
result.index = close.index

# result.to_csv(os.path.join(data_dir, 'result_alpha001_numpy.csv'))
print(f"---total time---: {t4 - t1}")